package com.etl.dataflow.executor.core;

import cn.hutool.core.util.StrUtil;
import com.etl.dataflow.common.entity.JdbcProperties;
import com.etl.dataflow.common.entity.TaskStatisticInfo;
import com.etl.dataflow.common.exception.EtlException;
import com.etl.dataflow.executor.element.*;
import com.etl.dataflow.executor.exception.UnsupportedFieldTypeException;
import com.etl.dataflow.executor.util.JdbcQuerySaveUtil;
import com.etl.dataflow.executor.util.Lob2JavaUtil;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.charset.Charset;
import java.sql.*;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author dx
 * @since 2022/4/14
 */
@Slf4j
public class OracleReader implements RdbmsReader {
    /**
     * 速度测量时间间隔(s)
     */
    private static final int SPEED_MEASURE_INTERVAL = 5;
    private final JdbcProperties jdbcProperties;
    private final String srcTableName;
    private final ChannelWriter channelWriter;
    private final AtomicBoolean reading = new AtomicBoolean(false);
    private final AtomicBoolean exit = new AtomicBoolean(false);
    private final TaskContext taskContext;
    private String sql = null;
    private String sqlForCount = null;
    private ResultSet resultSet;
    private Connection connection;
    private OnReadErrorListener onReadErrorListener;

    private AbstractAfterReadInterceptor afterReadInterceptor;

    public OracleReader(String srcTableName, JdbcProperties jdbcProperties, ChannelWriter channelWriter) {
        this.srcTableName = srcTableName;
        this.jdbcProperties = jdbcProperties;
        this.channelWriter = channelWriter;

        Optional<TaskContext> taskContextOptional = TaskContextHolder.getTaskContext();
        assert taskContextOptional.isPresent();
        taskContext = taskContextOptional.get();
    }

    private static final Set<String> RESERVED_WORD = new HashSet<>(Arrays.asList("level", "create", "alter", "table",
            "prior", "order", "by", "insert", "delete", "update", "index"));
    private static Column<?> createColumn(ResultSet resultSet, int index) throws SQLException, IOException {
        ResultSetMetaData metaData = resultSet.getMetaData();
        String columnName = metaData.getColumnLabel(index);
        if (RESERVED_WORD.contains(columnName)) {
            columnName = "\"" + columnName + "\"";
        }
        int columnType = metaData.getColumnType(index);
        String columnTypeName = metaData.getColumnTypeName(index);
        switch (columnType) {
            case Types.DATE:
                return new DateColumn(columnName, resultSet.getDate(index), columnType, columnTypeName);
            case Types.TIMESTAMP:
                return new TimestampColumn(columnName, resultSet.getTimestamp(index), columnType, columnTypeName);
            case Types.VARCHAR:
            case Types.NVARCHAR:
            case Types.CHAR:
            case Types.NCHAR:
                return new StringColumn(columnName, resultSet.getString(index), columnType, columnTypeName);
            case Types.CLOB:
                return new StringColumn(columnName, Lob2JavaUtil.clobToString(resultSet.getClob(index)), columnType, columnTypeName);
            case Types.BLOB:
                return new BytesColumn(columnName, Lob2JavaUtil.blobToBytes(resultSet.getBlob(index)), columnType, columnTypeName, Charset.forName("GBK"));
            case Types.NUMERIC:
            case Types.INTEGER:
            case Types.TINYINT:
            case Types.SMALLINT:
            case Types.BIGINT:
            case Types.DECIMAL:
            case Types.DOUBLE:
            case Types.FLOAT:
                return new BigDecimalColumn(columnName, resultSet.getBigDecimal(index), columnType, columnTypeName);
            default:
                throw new UnsupportedFieldTypeException();
        }
    }

    @Override
    public boolean startRead() {
        if (StrUtil.isBlank(srcTableName)) {
            throw new EtlException("请设置tableName");
        }
        try {
            boolean isClosed = resultSet != null && resultSet.isClosed();
            if (reading.get() || isClosed) {
                return false;
            }
            reading.compareAndSet(false, true);
            executeSqlAndWriteRowToChannel();
            return true;
        } catch (Exception e) {
            if (onReadErrorListener != null) {
                onReadErrorListener.onError(e);
            }
            throw new EtlException(e);
        } finally {
            TaskContextHolder.removeTaskContext();
        }
    }

    private void executeSqlAndWriteRowToChannel() throws SQLException, IOException {
        try {
            connection = JdbcQuerySaveUtil.getConnection(jdbcProperties);
            this.getRowCount();
            resultSet = JdbcQuerySaveUtil.query(connection, this.sql);
            long lastMeasureTime = System.currentTimeMillis();
            double interval;
            long rowNumber = 0L;
            long byteSize = 0L;
            while (resultSet.next() && !this.exit.get()) {
                int colCount = resultSet.getMetaData().getColumnCount();
                Row row = new DefaultRow();
                for (int i = 1; i <= colCount; i++) {
                    Column<?> column = createColumn(resultSet, i);
                    row.setColumn(i, column);
                }
                // 执行读后拦截器
                this.invokeAfterReadInterceptor(row);
                // 阻塞式写缓存
                channelWriter.writeOneRow(row);
                // 已读数量加1
                taskContext.getTaskStatisticInfo().getRedRowNumber().incrementAndGet();

                rowNumber++;
                byteSize += row.getContentByteSize();

                interval = (System.currentTimeMillis() - lastMeasureTime) / 1e+3;
                if (interval >= SPEED_MEASURE_INTERVAL) {
                    taskContext.getTaskStatisticInfo().getWriteBytesPerSecond().set((long) (byteSize / interval));
                    byteSize = 0;
                    taskContext.getTaskStatisticInfo().getWriteRowsPerSecond().set((long) (rowNumber / interval));
                    rowNumber = 0;
                    lastMeasureTime = System.currentTimeMillis();
                }
            }
        } finally {
            log.debug("读取结束。。。。");
            JdbcQuerySaveUtil.closeResultSet(resultSet);
            reading.compareAndSet(true, false);
            exit();
        }
    }

    private void invokeAfterReadInterceptor(Row row) {
        if (this.afterReadInterceptor != null) {
            this.afterReadInterceptor.onAfterRead(row);
        }
    }

    @Override
    public void exit() {
        this.exit.set(true);
        this.channelWriter.terminate();
    }

    @Override
    public boolean setSql(String sql) {
        if (this.reading.get()) {
            return false;
        }
        this.sql = sql;
        this.sqlForCount = String.format("select count(*) from (%s)", sql);
        return true;
    }

    public String getSql() {
        return sql;
    }

    @Override
    public long getRowCount() {
        Optional<TaskContext> taskContext = TaskContextHolder.getTaskContext();
        if (!taskContext.isPresent()) {
            return -1;
        }
        TaskStatisticInfo statisticInfo = taskContext.get().getTaskStatisticInfo();
        if (statisticInfo.getSrcRowNumber().get() != -1) {
            return statisticInfo.getSrcRowNumber().get();
        }
        try {
            statisticInfo.getSrcRowNumber().set(
                    JdbcQuerySaveUtil.count(connection, this.sqlForCount));
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return statisticInfo.getSrcRowNumber().get();
    }

    @Override
    public long getRedRowNumber() {
        return this.taskContext.getTaskStatisticInfo().getRedRowNumber().get();
    }

    @Override
    public void setOnReadErrorListener(OnReadErrorListener listener) {
        this.onReadErrorListener = listener;
    }

    @Override
    public void setAfterReadInterceptor(AbstractAfterReadInterceptor interceptor) {
        this.afterReadInterceptor = interceptor;
    }

}
